from __future__ import annotations

import logging
import struct
from copy import deepcopy
from dataclasses import dataclass
from typing import IO, List

import numpy as np

riff_wave_logger = logging.getLogger('RiffWaveIOLogger')
formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] - %(message)s')
handler = logging.StreamHandler()
handler.setFormatter(formatter)

riff_wave_logger.addHandler(handler)
riff_wave_logger.setLevel(logging.INFO)


@dataclass()
class SimpleChunk:
    type: str
    size: int
    content_bytes: bytes

    @classmethod
    def from_bytes(cls, data: bytes) -> SimpleChunk:
        return SimpleChunk(data[:4].decode('ascii'), len(data) - 8, data[8:])

    def dump(self) -> bytes:
        header_bytes = struct.pack('<4sI', self.type.encode('ascii'), self.size)
        return header_bytes + self.content_bytes


@dataclass()
class WavFormat:
    """
    format section 共计 16 + 4(fmt ) + 4(size) = 20 字节
        12 ~ 16 字节固定为："fmt "
        16 ~ 19 4字节 fmt 的有效内容 size，固定为 16，即：10 00 00 00
        20 ~ 21 2字节为表示编码格式 1 为 pcm 编码
        22 ~ 23 2字节为声道数, 1 为单声道, 2 为立体声
        24 ~ 27 4字节为采样率
        28 ~ 31 4字节为每秒字节数，即比特率
        32 ~ 33 2字节为数据块长度（通道数 x 位宽）
        34 ~ 35 2字节为采样位数
    """
    encoding_code: int
    n_channel: int
    sample_rate: int
    byte_rate: int
    chunk_bytes: int = 2
    sample_bits: int = 8

    @property
    def sample_width(self):
        return self.byte_rate // self.sample_rate // 2

    @classmethod
    def from_bytes(cls, data: bytes) -> WavFormat:
        if len(data) != 24:
            raise ValueError('No enough bytes for parsing fmt.')

        if data[:4] != b'fmt ':
            raise ValueError(f'Not valid fmt bytes: {data}')

        (
            encoding_code,
            n_channels,
            sample_rate,
            byte_rate,
            chunk_bytes,
            sample_bits
        ) = struct.unpack('<HHIIHH', data[8:])

        return cls(encoding_code, n_channels, sample_rate, byte_rate, chunk_bytes, sample_bits)

    def dump(self) -> bytes:
        return struct.pack(
            '<4sIHHIIHH',
            b'fmt ',
            16,
            self.encoding_code,
            self.n_channel,
            self.sample_rate,
            self.byte_rate,
            self.chunk_bytes,
            self.sample_bits
        )


@dataclass()
class WaveData:
    """
        仅支持解析 pcm 格式的 wav 文件
    """
    bytes_size: int
    samples: np.ndarray  # (n_samples) or (n_samples, n_channels>1)
    fmt: WavFormat

    @classmethod
    def from_bytes(cls, data: bytes, format_: WavFormat) -> WaveData:
        _, bytes_size = struct.unpack('<4sI', data[:8])
        riff_wave_logger.debug(f'wave data bytes_size: {bytes_size}')
        sample_width = format_.sample_width
        if sample_width == 2:
            dtype = np.int16
        elif sample_width == 3:
            dtype = np.dtype([('int24', 'i4')])  # 24-bit int stored in 32-bit container
        elif sample_width == 4:
            dtype = np.int32
        else:
            raise ValueError(f"Unsupported sample width: {sample_width}")

        samples = np.frombuffer(data[8:], dtype=dtype)
        if sample_width == 3:
            samples = samples['int24'] >> 8  # Adjust 24-bit int to correct range

        if format_.n_channel > 1:
            samples = samples.reshape(-1, format_.n_channel)

        # 归一化float
        if dtype == np.int16:
            samples = samples / 2 ** 15
        elif dtype == np.int32:
            samples = samples / 2 ** 31

        return cls(bytes_size, samples, deepcopy(format_))

    def dump(self) -> bytes:
        # samples to int
        sample_width = self.fmt.sample_width
        if sample_width == 2:
            dtype = np.int16
        elif sample_width == 3:
            dtype = np.dtype([('int24', 'i4')])
        elif sample_width == 4:
            dtype = np.int32
        else:
            raise ValueError(f"Unsupported sample width: {sample_width}")

        sample_dtype = self.samples.dtype
        if dtype == np.int16:
            samples = self.samples * 2 ** 15
        elif dtype == np.int32:
            samples = self.samples * 2 ** 31
        else:
            raise ValueError(f"Unsupported dtype: {sample_dtype}")

        samples = samples.astype(dtype)
        if sample_width == 3:
            samples = samples << 8

        if self.fmt.n_channel > 1:
            samples = samples.flatten()

        header_size = struct.pack('<4sI', b'data', self.bytes_size)
        return header_size + samples.tobytes()


@dataclass()
class AdtlLabel:
    size: int = 0
    id: int = 0
    text: str = ''

    @classmethod
    def from_bytes(cls, data: bytes) -> AdtlLabel:
        if len(data) < 12:
            raise ValueError('No enough bytes for label.')

        _, size, idx = struct.unpack('<4sII', data[:12])
        text = data[12:].decode('utf-8').strip('\x00')
        return AdtlLabel(size, idx, text)

    @classmethod
    def from_text(cls, text: str, idx: int) -> AdtlLabel:
        text_bytes = text.encode('utf-8')
        size = len(text_bytes) + (2 - len(text_bytes) % 2) + 4
        return cls(size=size, id=idx, text=text)

    def dump(self) -> bytes:
        content_bytes = self.text.encode('utf-8')
        header_bytes = struct.pack('<4sII', b'labl', len(content_bytes) + 4, self.id)
        content_bytes += b'\x00' * (2 - len(content_bytes) % 2) # 对齐到字节并补充结束符
        return header_bytes + content_bytes


@dataclass()
class AdtlLabelList:
    labels: List[AdtlLabel] = None

    def __post_init__(self):
        self.labels = self.labels or []

    @classmethod
    def from_bytes(cls, data: bytes) -> AdtlLabelList:
        _, bytes_size, adtl = struct.unpack('<4sI4s', data[:12])
        label_list = cls()
        i = 12
        while i < len(data):
            labl, label_size = struct.unpack('<4sI', data[i:i + 8])
            cur_label = AdtlLabel.from_bytes(data[i:i + label_size + 8])
            i += label_size + 8
            label_list.labels.append(cur_label)

        return label_list

    def dump(self) -> bytes:
        return struct.pack('<4sI4s', b'LIST', sum([label.size + 8 for label in self.labels]) + 4, b'adtl') + \
               b''.join([label.dump() for label in self.labels])


@dataclass()
class CueTiming:
    id: int = 0
    start: int = 0
    end: int = 0

    @classmethod
    def from_bytes(cls, data: bytes) -> CueTiming:
        if len(data) != 24:
            raise ValueError(f"Invalid cue timing data: {data}")

        id_, start, _, _, _, end = struct.unpack('<II4sIII', data)
        return cls(id_, start, end)

    def dump(self) -> bytes:
        return struct.pack('<II4sIII', self.id, self.start, b'data', 0, 0, self.end)


@dataclass()
class Cue:
    timings: List[CueTiming] = None

    def __post_init__(self):
        self.timings = self.timings or []

    @property
    def bytes_size(self) -> int:
        return sum([24 for _ in self.timings]) + 4

    @classmethod
    def from_bytes(cls, data: bytes) -> Cue:
        cue_str, bytes_size, seq_num, = struct.unpack('<4sII', data[:12])
        timings: List[CueTiming] = []
        for i in range(12, len(data), 24):
            cur_timing = CueTiming.from_bytes(data[i:i + 24])
            timings.append(cur_timing)

        cue = cls(timings=timings)

        return cue

    def dump(self) -> bytes:
        header_bytes = struct.pack('<4sII', b'cue ', self.bytes_size, len(self.timings))
        content_bytes = b''.join([timing.dump() for timing in self.timings])
        return header_bytes + content_bytes

@dataclass()
class AdtlAnnotation:
    id: int
    start_idx: int
    end_idx: int
    label: str
    sample_rate: int

    @property
    def start_time(self) -> float:
        return self.start_idx / self.sample_rate

    @property
    def end_time(self) -> float:
        return self.end_idx / self.sample_rate


@dataclass()
class WaveContainer:
    """
        RIFF 格式的 wav 文件解析容器
    """
    io: IO
    bytes_size: int
    fmt: WavFormat
    data: WaveData = None
    sub_chunks: List = None

    def __post_init__(self):
        self.sub_chunks = self.sub_chunks or []

    @property
    def sample_rate(self) -> int:
        return self.fmt.sample_rate

    @property
    def duration(self) -> float:
        return len(self.data.samples) / self.fmt.sample_rate

    @property
    def wav_array(self) -> np.ndarray:
        return self.data.samples

    @classmethod
    def from_stream(cls, stream: IO) -> WaveContainer:
        header = stream.read(12)
        if len(header) < 12:
            raise ValueError('No enough bytes to parse.')

        riff, bytes_size, wave = struct.unpack('<4sI4s', header)
        if riff != b'RIFF':
            raise ValueError('Not a wav file.')

        if wave != b'WAVE':
            raise ValueError('Not a wav file.')

        # bytes_size is the length of valid bytes content after.
        fmt = WavFormat.from_bytes(stream.read(24))

        if fmt.encoding_code != 1:
            raise ValueError('Only support PCM encoding.')

        wave_riff = cls(
            io=stream,
            bytes_size=bytes_size,
            fmt=fmt,
        )

        if wave_riff is None:
            riff_wave_logger.warning('No wav samples found.')

        # 顺序地读取所有 sub-chunks
        while stream.tell() < wave_riff.bytes_size + 8:
            chunk_name, chunk_size = struct.unpack('<4sI', stream.read(8))
            if chunk_name == b'00000000':
                stream.seek(stream.tell() - 4)
                continue

            elif chunk_name == b'data':
                stream.seek(stream.tell() - 8)
                wave_riff.data = WaveData.from_bytes(stream.read(chunk_size + 8), wave_riff.fmt)\

            elif chunk_name == b'LIST':
                list_type = struct.unpack('<4s', stream.read(4))[0]
                if list_type == b'adtl':
                    stream.seek(stream.tell() - 12)
                    wave_riff.sub_chunks.append(AdtlLabelList.from_bytes(stream.read(chunk_size + 8)))
                else:
                    stream.seek(stream.tell() - 12)
                    wave_riff.sub_chunks.append(SimpleChunk.from_bytes(stream.read(chunk_size + 8)))

            elif chunk_name == b'cue ':
                stream.seek(stream.tell() - 8)
                wave_riff.sub_chunks.append(Cue.from_bytes(stream.read(chunk_size + 8)))

            else:
                stream.seek(stream.tell() - 8)
                riff_wave_logger.debug(f'unknown chunk: {chunk_name}, size: {chunk_size}')
                wave_riff.sub_chunks.append(SimpleChunk.from_bytes(stream.read(chunk_size + 8)))

        return wave_riff

    def delete_adtl_annotations(self):
        flag_has_adtl = any(isinstance(x, (Cue, AdtlLabelList)) for x in self.sub_chunks)
        if not flag_has_adtl:
            return

        self.sub_chunks = [chunk for chunk in self.sub_chunks if not isinstance(chunk, (Cue, AdtlLabelList))]

    def update_adtl_annotations(self, adtl_annotations: List[AdtlAnnotation]):

        cue = Cue()
        adtl_label_list = AdtlLabelList()

        cue.timings = [CueTiming(id=anno.id, start=anno.start_idx, end=anno.end_idx) for anno in adtl_annotations]
        adtl_label_list.labels = [AdtlLabel.from_text(anno.label, anno.id) for anno in adtl_annotations]

        self.delete_adtl_annotations()
        self.sub_chunks.insert(0, cue)
        self.sub_chunks.insert(1, adtl_label_list)

    def get_adtl_annotations(self) -> List[AdtlAnnotation]:
        sample_rate = self.sample_rate
        cue: Cue = None
        adtl_label_list: AdtlLabelList = None
        for chunk in self.sub_chunks:
            if isinstance(chunk, Cue):
                cue = chunk
            elif isinstance(chunk, AdtlLabelList):
                adtl_label_list = chunk

        if cue is None or adtl_label_list is None:
            return []

        adtl_annotations = []
        id2text = {label.id: label.text for label in adtl_label_list.labels}
        for timing in cue.timings:
            if timing.id not in id2text:
                continue

            adtl_annotations.append(AdtlAnnotation(
                id=timing.id,
                start_idx=timing.start,
                end_idx=timing.end,
                label=id2text[timing.id],
                sample_rate=sample_rate
            ))

        return adtl_annotations

    def dump(self) -> bytes:
        riff_bytes = struct.pack('<4sI4s', b'RIFF', self.bytes_size, b'WAVE') + self.fmt.dump()
        data_bytes = self.data.dump()
        sub_chunks_bytes = b''.join([chunk.dump() for chunk in self.sub_chunks])

        # update bytes_size
        real_bytes_size = len(riff_bytes + data_bytes + sub_chunks_bytes) - 8
        riff_bytes = struct.pack('<4sI4s', b'RIFF', real_bytes_size, b'WAVE') + self.fmt.dump()
        return riff_bytes + data_bytes + sub_chunks_bytes
